import paramiko
import sys
paramiko.util.log_to_file('paramiko.log')


class SSHService:
    def __init__(self, ipaddr, username="root", password=None, pkey_path=None):
        self.__sshclient = paramiko.SSHClient()
        self.__sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.__ipaddr = ipaddr
        self.__username = username
        self.__password = password
        self.__pkey_path = pkey_path
        self.__setup_ssh_connection()

    def __setup_ssh_connection(self):
        if self.__password:
            self.__sshclient.connect(hostname=self.__ipaddr, username=self.__username, password=self.__password)
        elif self.__pkey_path:
            pkey = paramiko.RSAKey.from_private_key_file(self.__pkey_path)
            self.__sshclient.connect(hostname=self.__ipaddr, username=self.__username, pkey=pkey)
        else:
            print("ssh connection setup fail!")

    def run_cmd(self, cmd):
        _, stdout, stderr = self.__sshclient.exec_command(cmd)
        error_info = stderr.read()
        if error_info:
            print(error_info)
        else:
            return stdout.read()

    def close_ssh(self):
        self.__sshclient.close()
        